#include "core/main.h"

class SarsaLearnerParameters
{
public:
    SarsaLearnerParameters()
    {
        ////
        //// Frogger (Hard)
        ////
        //REWARD_SCALE = 1000.0;

        //EPSILON_0 = 0.25;  // Exploration rate
        //GAMMA_0   = 0.9;  // Discount factor
        //ALPHA_0   = 0.01; // Learning rate
        //LAMBDA_0  = 0.2;  // Eligibility discount

        //EPSILON_HALF_LIFE_FRAMES = 10000;
        //GAMMA_HALF_LIFE_FRAMES   = -1;
        //ALPHA_HALF_LIFE_FRAMES   = -1;
        //LAMBDA_HALF_LIFE_FRAMES  = -1;

        //EPSILON   = EPSILON_0; // Exploration rate
        //GAMMA     = GAMMA_0;   // Discount factor
        //ALPHA     = ALPHA_0;   // Learning rate
        //LAMBDA    = LAMBDA_0;  // Eligibility discount

        //NUM_FRAMES_ELAPSED = 0;

        ////
        //// Pong
        ////
        //REWARD_SCALE = 1000.0;

        //EPSILON_0 = 0.5;   // Exploration rate
        //GAMMA_0   = 0.9;   // Discount factor
        //ALPHA_0   = 0.001; // Learning rate
        //LAMBDA_0  = 0.999;  // Eligibility discount

        //EPSILON_HALF_LIFE_FRAMES = 50000;
        //GAMMA_HALF_LIFE_FRAMES   = -1;
        //ALPHA_HALF_LIFE_FRAMES   = -1;
        //LAMBDA_HALF_LIFE_FRAMES  = -1;

        //EPSILON   = EPSILON_0; // Exploration rate
        //GAMMA     = GAMMA_0;   // Discount factor
        //ALPHA     = ALPHA_0;   // Learning rate
        //LAMBDA    = LAMBDA_0;  // Eligibility discount

        //NUM_FRAMES_ELAPSED = 0;

        REWARD_SCALE = 1000.0;

        EPSILON_0 = 0.5; // Exploration rate
        GAMMA_0   = 0.9;  // Discount factor
        ALPHA_0   = 0.1;  // Learning rate
        LAMBDA_0  = 0.9;  // Eligibility discount

        EPSILON_HALF_LIFE_FRAMES = 20000;
        GAMMA_HALF_LIFE_FRAMES   = -1;
        ALPHA_HALF_LIFE_FRAMES   = -1;
        LAMBDA_HALF_LIFE_FRAMES  = -1;

        EPSILON   = EPSILON_0; // Exploration rate
        GAMMA     = GAMMA_0;   // Discount factor
        ALPHA     = ALPHA_0;   // Learning rate
        LAMBDA    = LAMBDA_0;  // Eligibility discount

        NUM_FRAMES_ELAPSED = 0;
    }

    void Update()
    {
        NUM_FRAMES_ELAPSED++;

        if ( EPSILON_HALF_LIFE_FRAMES > 0 )
        {
            EPSILON = EPSILON_0 * pow( 2.0, double(-NUM_FRAMES_ELAPSED) / double(EPSILON_HALF_LIFE_FRAMES) );
        }

        if ( GAMMA_HALF_LIFE_FRAMES > 0 )
        {
            GAMMA = GAMMA_0 * pow( 2.0, -NUM_FRAMES_ELAPSED / GAMMA_HALF_LIFE_FRAMES );
        }

        if ( ALPHA_HALF_LIFE_FRAMES > 0 )
        {
            ALPHA = ALPHA_0 * pow( 2.0, -NUM_FRAMES_ELAPSED / ALPHA_HALF_LIFE_FRAMES );
        }

        if ( LAMBDA_HALF_LIFE_FRAMES > 0 )
        {
            LAMBDA = LAMBDA_0 * pow( 2.0, -NUM_FRAMES_ELAPSED / LAMBDA_HALF_LIFE_FRAMES );
        }
    }

    void Print()
    {
        std::stringstream ss;
        ss << EPSILON;
        mike::GlTextConsolePrint( "EPSILON", ss.str() );
        ss.str("");

        ss << GAMMA;
        mike::GlTextConsolePrint( "GAMMA", ss.str() );
        ss.str("");

        ss << ALPHA;
        mike::GlTextConsolePrint( "ALPHA", ss.str() );
        ss.str("");

        ss << LAMBDA;
        mike::GlTextConsolePrint( "LAMBDA", ss.str() );
        ss.str("");
    }

    double REWARD_SCALE;

    double EPSILON_0; // Exploration rate
    double GAMMA_0;   // Discount factor
    double ALPHA_0;   // Learning rate
    double LAMBDA_0;  // Eligibility discount

    int EPSILON_HALF_LIFE_FRAMES;
    int GAMMA_HALF_LIFE_FRAMES;
    int ALPHA_HALF_LIFE_FRAMES;
    int LAMBDA_HALF_LIFE_FRAMES;

    double EPSILON; // Exploration rate
    double GAMMA;   // Discount factor
    double ALPHA;   // Learning rate
    double LAMBDA;  // Eligibility discount

    int NUM_FRAMES_ELAPSED;
};

static SarsaLearnerParameters sSarsaLearnerParameters;

SarsaLearner::SarsaLearner(Game *game, const SARSA_TRANSFORM_TYPE transform, const vector<ControllerState> &actions)
    : GameAgent( game ), _transform(transform), _actions(actions)
{
    // First figure out the dimensionality of the feature vector for this game
    // vector<double> featureVector;
    // _transform.transform(_game->ActiveFrame(), 0, featureVector);
    _featureDim = _transform.dimension();

    // Initialize the weight vector
   _weightVector.resize(_featureDim, 0.2);

#ifdef USE_DENSE_TRACE
    _traceVector = new DenseDoubleVector(_featureDim);
#else
    _traceVector = new SparseDoubleVector();
#endif

   // Pick the first action
   _nextAction = pickAction(sSarsaLearnerParameters.EPSILON);
}

void SarsaLearner::takeStep()
{   
    // Copy the current frame and action before taking the action
    Frame curState = _game->ActiveFrame();
    int curAction = _nextAction;

    // double startTime = mike::TimeGetTimeDeltaSeconds();

    // Take the next action (that we have already picked)
    _game->Step(_actions[curAction]);

    Frame nextState = _game->ActiveFrame();
    double reward = getReward(nextState) / sSarsaLearnerParameters.REWARD_SCALE;

    stringstream ss;
    ss << reward;
    mike::GlTextConsolePrint("Reward", ss.str());
    ss.str("");

    int nextAction = pickAction(sSarsaLearnerParameters.EPSILON);
    double delta = reward;
    delta += sSarsaLearnerParameters.GAMMA * getValue(nextState, nextAction);
    delta -= getValue(curState, curAction);

    ss << delta;
    mike::GlTextConsolePrint("delta", ss.str());
    mike::MatrixWriterAppendDouble("delta", delta);
    ss.str("");

    // SparseIntVector curFeature;
#ifdef USE_DENSE_FEATURES
    DenseIntVector curFeature(_featureDim);
#else
    SparseIntVector curFeature;
#endif
    _transform.transform(curState, curAction, curFeature);

    // Update the trace vector and the weight vector
    *_traceVector *= sSarsaLearnerParameters.LAMBDA * sSarsaLearnerParameters.GAMMA;
    for (auto it = curFeature.indices().begin(); it != curFeature.indices().end(); ++it) {
        _traceVector->set(*it, 1);
    }
    
    ss << _traceVector->numNonZero() << " / " << _featureDim;
    mike::GlTextConsolePrint( "Active traces", ss.str() );
    ss.str("");

    _traceVector->filter(sSarsaLearnerParameters.ALPHA * delta, _weightVector);

    _nextAction = nextAction;

    sSarsaLearnerParameters.Update();
    sSarsaLearnerParameters.Print();

    // double endTime = mike::TimeGetTimeDeltaSeconds();
    // _computeTimes.push_back(endTime);
}

double SarsaLearner::getTotalComputeTime()
{
    sort(_computeTimes.begin(), _computeTimes.end());
    double sum = 0;
    for (auto it = _computeTimes.begin(); it != _computeTimes.end(); ++it) sum += *it;
    return sum;
}

// Pick an action to take from the current frame of the game using an
// epsilon-greedy policy. Returns the index of the next action to take.
int SarsaLearner::pickAction(double epsilon)
{
    int action;

    if (rnd() < epsilon) {
        // Take a random action
        action = rand() % _actions.size();
        // action = 4; // Go up
    } else {
        action = 0;
        const Frame &curFrame = _game->ActiveFrame();
        double bestValue = getValue(curFrame, 0);
        double secondBestValue = -DBL_MAX;

        for (UINT i = 1; i < _actions.size(); i++) {
            double value = getValue(curFrame, i);
            if (value > bestValue) {
                secondBestValue = bestValue;
                bestValue = value;
                action = i;
            }
            if (value > secondBestValue && value < bestValue) {
                secondBestValue = value;
            }
        }
    }

    return action;
}

double SarsaLearner::getValue(const Frame &frame, int actionIndex)
{
    // SparseIntVector featureVector;
#ifdef USE_DENSE_FEATURES
    DenseIntVector featureVector(_featureDim);
#else
    SparseIntVector featureVector;
#endif
    _transform.transform(frame, actionIndex, featureVector);

    return featureVector.dot(_weightVector);
}

double SarsaLearner::getReward(const Frame &frame)
{
    return frame.observedInts[2];
}